from pydantic import BaseModel
from typing import List, Union


# 基础数据结构
class XAxis(BaseModel):
    label: str
    data: List[Union[int, float, str]]


class YAxis(BaseModel):
    label: str
    data: List[Union[int, float]]


class Series(BaseModel):
    name: str
    data: List[Union[List[Union[int, float]], Union[int, float]]]
    type: str
    smooth: bool = False  # 只针对折线图适用，默认不平滑


# 针对折线图 (Line Chart)
class LineChart(BaseModel):
    chartType: str = "line"
    xAxis: XAxis
    yAxis: YAxis
    series: List[Series]


# 针对直方图 (Histogram)
class BarChart(BaseModel):
    chartType: str = "bar"
    xAxis: XAxis
    yAxis: YAxis
    series: List[Series]


# 针对散点图 (Scatter Plot)
class ScatterPlot(BaseModel):
    chartType: str = "scatter"
    xAxis: XAxis
    yAxis: YAxis
    series: List[Series]


# 综合模型，支持所有类型的图表
class Chart(BaseModel):
    chartType: str
    xAxis: XAxis
    yAxis: YAxis
    series: List[Series]


# 示例：折线图的 Pydantic 模型
line_chart = LineChart(
    xAxis=XAxis(label="Epoch", data=[1, 2, 3, 4, 5]),
    yAxis=YAxis(label="Loss", data=[0.5, 0.4, 0.35, 0.3, 0.25]),
    series=[Series(name="Training Loss", data=[0.5, 0.4, 0.35, 0.3, 0.25], type="line", smooth=True)]
)

# 示例：直方图的 Pydantic 模型
bar_chart = BarChart(
    xAxis=XAxis(label="Categories", data=["Cat1", "Cat2", "Cat3", "Cat4"]),
    yAxis=YAxis(label="Frequency", data=[10, 20, 30, 40]),
    series=[Series(name="Category Frequency", data=[10, 20, 30, 40], type="bar")]
)

# 示例：散点图的 Pydantic 模型
scatter_plot = ScatterPlot(
    xAxis=XAxis(label="Parameter A", data=[0.1, 0.2, 0.3, 0.4, 0.5]),
    yAxis=YAxis(label="Loss", data=[0.5, 0.4, 0.35, 0.3, 0.25]),
    series=[Series(name="Loss vs Parameter A", data=[[0.1, 0.5], [0.2, 0.4], [0.3, 0.35], [0.4, 0.3], [0.5, 0.25]],
                   type="scatter")]
)
